#include "compi/context.h"
#include "compi/resource_manager.h"

#include <algorithm>
#include <functional>
#include <mpi.h>
#include <stdexcept>

namespace compi
{

// 静态成员定义
std::unordered_map<std::size_t, std::unique_ptr<ResourceManager::ResourceEntry>>
    ResourceManager::s_resources;
std::mutex ResourceManager::s_resources_mutex;

void
ResourceManager::TagAllocatorDeleter::operator()(MessageTagAllocator * ptr) const
{
    delete ptr;
}

std::size_t
ResourceManager::compute_comm_hash(MPI_Comm comm)
{
    MPI_Fint handle = MPI_Comm_c2f(comm);
    std::hash<MPI_Fint> hasher;
    return hasher(handle);
}

ResourceManager::ResourceEntry *
ResourceManager::find_resource_entry(MPI_Comm comm)
{
    std::size_t hash = compute_comm_hash(comm);
    auto it = s_resources.find(hash);

    if (it != s_resources.end())
    {
        // 验证communicator是否真的相同
        int result;
        if (MPI_Comm_compare(comm, it->second->comm, &result) == MPI_SUCCESS)
        {
            if (result == MPI_IDENT || result == MPI_CONGRUENT)
            {
                return it->second.get();
            }
        }

        // 哈希冲突，需要移除无效条目
        s_resources.erase(it);
    }

    return nullptr;
}

std::unique_ptr<ResourceManager::ResourceEntry>
ResourceManager::create_resource_entry(MPI_Comm comm)
{
    auto entry = std::make_unique<ResourceEntry>(comm);

    // 创建TagAllocator实例，使用自定义删除器
    entry->tag_allocator = std::unique_ptr<MessageTagAllocator, TagAllocatorDeleter>(
        new MessageTagAllocator(comm), TagAllocatorDeleter{});

    return entry;
}

MessageTagAllocator &
ResourceManager::get_tag_allocator(MPI_Comm comm)
{
    std::lock_guard<std::mutex> lock(s_resources_mutex);

    // 查找现有条目
    ResourceEntry * entry = find_resource_entry(comm);
    if (entry != nullptr)
    {
        return *entry->tag_allocator;
    }

    // 创建新条目
    auto new_entry = create_resource_entry(comm);
    MessageTagAllocator * allocator = new_entry->tag_allocator.get();

    std::size_t hash = compute_comm_hash(comm);
    s_resources[hash] = std::move(new_entry);

    return *allocator;
}

void
ResourceManager::destroy_tag_allocator(MPI_Comm comm)
{
    std::lock_guard<std::mutex> lock(s_resources_mutex);

    std::size_t hash = compute_comm_hash(comm);
    auto it = s_resources.find(hash);

    if (it != s_resources.end())
    {
        // 验证communicator是否真的相同
        int result;
        if (MPI_Comm_compare(comm, it->second->comm, &result) == MPI_SUCCESS)
        {
            if (result == MPI_IDENT || result == MPI_CONGRUENT)
            {
                s_resources.erase(it);
                return;
            }
        }
    }
}

void
ResourceManager::cleanup_all()
{
    std::lock_guard<std::mutex> lock(s_resources_mutex);
    s_resources.clear();
}

ResourceManager::ResourceStats
ResourceManager::get_stats()
{
    std::lock_guard<std::mutex> lock(s_resources_mutex);

    ResourceStats stats{};
    stats.active_allocators = static_cast<int>(s_resources.size());

    for (const auto & [hash, entry] : s_resources)
    {
        if (entry && entry->tag_allocator)
        {
            auto global_stats = entry->tag_allocator->get_global_stats();
            stats.total_namespaces += global_stats.total_namespaces;
            stats.total_allocated_tags += global_stats.total_allocated_tags;
        }
    }

    return stats;
}

} // namespace compi
